from sentence_transformers import SentenceTransformer
from langchain.embeddings.base import Embeddings
from typing import List
import yaml
import numpy as np


class LocalEmbeddings(Embeddings):
    """本地embedding模型封装"""

    def __init__(self, config_path: str = "config.yaml"):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)

        self.embed_config = config['embeddings']
        self.model = SentenceTransformer(
            self.embed_config['model_path'],
            device=self.embed_config['device']
        )

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """嵌入文档列表"""
        embeddings = self.model.encode(
            texts,
            batch_size=self.embed_config['batch_size'],
            show_progress_bar=True
        )
        return embeddings.tolist()

    def embed_query(self, text: str) -> List[float]:
        """嵌入单个查询"""
        embedding = self.model.encode([text])
        return embedding[0].tolist()